{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "28S76DVlfCMZ"
      },
      "source": [
        "# Adversarial Reprogramming of Growing Neural Cellular Automata\n",
        "\n",
        "This notebook contains code to reproduce experiments and figures regarding Growing CAs for the \"Adversarial Reprogramming of Neural Cellular Automata\" article.\n",
        "\n",
        "Copyright 2020 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "REhz9l8Aj7W0"
      },
      "source": [
        "# Common code for both experiments"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "i5wi_r4gyzFr"
      },
      "outputs": [],
      "source": [
        "#@title Imports and Notebook Utilities\n",
        "%tensorflow_version 2.x\n",
        "\n",
        "import os\n",
        "import io\n",
        "import PIL.Image, PIL.ImageDraw\n",
        "import base64\n",
        "import zipfile\n",
        "import json\n",
        "import requests\n",
        "import numpy as np\n",
        "import matplotlib.pylab as pl\n",
        "import glob\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from IPython.display import Image, HTML, clear_output\n",
        "import tqdm\n",
        "\n",
        "import os\n",
        "os.environ['FFMPEG_BINARY'] = 'ffmpeg'\n",
        "import moviepy.editor as mvp\n",
        "from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter\n",
        "clear_output()\n",
        "\n",
        "def np2pil(a):\n",
        "  if a.dtype in [np.float32, np.float64]:\n",
        "    a = np.uint8(np.clip(a, 0, 1)*255)\n",
        "  return PIL.Image.fromarray(a)\n",
        "\n",
        "def imwrite(f, a, fmt=None):\n",
        "  a = np.asarray(a)\n",
        "  if isinstance(f, str):\n",
        "    fmt = f.rsplit('.', 1)[-1].lower()\n",
        "    if fmt == 'jpg':\n",
        "      fmt = 'jpeg'\n",
        "    f = open(f, 'wb')\n",
        "  np2pil(a).save(f, fmt, quality=95)\n",
        "\n",
        "def imencode(a, fmt='jpeg'):\n",
        "  a = np.asarray(a)\n",
        "  if len(a.shape) == 3 and a.shape[-1] == 4:\n",
        "    fmt = 'png'\n",
        "  f = io.BytesIO()\n",
        "  imwrite(f, a, fmt)\n",
        "  return f.getvalue()\n",
        "\n",
        "def im2url(a, fmt='jpeg'):\n",
        "  encoded = imencode(a, fmt)\n",
        "  base64_byte_string = base64.b64encode(encoded).decode('ascii')\n",
        "  return 'data:image/' + fmt.upper() + ';base64,' + base64_byte_string\n",
        "\n",
        "def imshow(a, fmt='jpeg'):\n",
        "  display(Image(data=imencode(a, fmt)))\n",
        "\n",
        "def tile2d(a, w=None):\n",
        "  a = np.asarray(a)\n",
        "  if w is None:\n",
        "    w = int(np.ceil(np.sqrt(len(a))))\n",
        "  th, tw = a.shape[1:3]\n",
        "  pad = (w-len(a))%w\n",
        "  a = np.pad(a, [(0, pad)]+[(0, 0)]*(a.ndim-1), 'constant')\n",
        "  h = len(a)//w\n",
        "  a = a.reshape([h, w]+list(a.shape[1:]))\n",
        "  a = np.rollaxis(a, 2, 1).reshape([th*h, tw*w]+list(a.shape[4:]))\n",
        "  return a\n",
        "\n",
        "def zoom(img, scale=4):\n",
        "  img = np.repeat(img, scale, 0)\n",
        "  img = np.repeat(img, scale, 1)\n",
        "  return img\n",
        "\n",
        "class VideoWriter:\n",
        "  def __init__(self, filename, fps=30.0, **kw):\n",
        "    self.writer = None\n",
        "    self.params = dict(filename=filename, fps=fps, **kw)\n",
        "\n",
        "  def add(self, img):\n",
        "    img = np.asarray(img)\n",
        "    if self.writer is None:\n",
        "      h, w = img.shape[:2]\n",
        "      self.writer = FFMPEG_VideoWriter(size=(w, h), **self.params)\n",
        "    if img.dtype in [np.float32, np.float64]:\n",
        "      img = np.uint8(img.clip(0, 1)*255)\n",
        "    if len(img.shape) == 2:\n",
        "      img = np.repeat(img[..., None], 3, -1)\n",
        "    self.writer.write_frame(img)\n",
        "\n",
        "  def close(self):\n",
        "    if self.writer:\n",
        "      self.writer.close()\n",
        "\n",
        "  def __enter__(self):\n",
        "    return self\n",
        "\n",
        "  def __exit__(self, *kw):\n",
        "    self.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zR6I1JONmWBb"
      },
      "outputs": [],
      "source": [
        "#@title Cellular Automata Parameters\n",
        "CHANNEL_N = 16        # Number of CA state channels\n",
        "TARGET_SIZE = 40\n",
        "BATCH_SIZE = 8\n",
        "POOL_SIZE = 1024\n",
        "CELL_FIRE_RATE = 0.5\n",
        "\n",
        "EXPERIMENT_TYPE = \"Regenerating\" #@param [\"Persistent\", \"Regenerating\"]\n",
        "EXPERIMENT_MAP = {\"Persistent\":0, \"Regenerating\":1}\n",
        "EXPERIMENT_N = EXPERIMENT_MAP[EXPERIMENT_TYPE]\n",
        "\n",
        "DAMAGE_N = [0, 3][EXPERIMENT_N]  # Number of patterns to damage in a batch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lCbPFbI_zosW"
      },
      "outputs": [],
      "source": [
        "#@title CA Model and Utilities\n",
        "#@markdown This model doesn't have a bias in the last layer.\n",
        "from tensorflow.keras.layers import Conv2D\n",
        "\n",
        "def load_image(url, max_size=TARGET_SIZE):\n",
        "  r = requests.get(url)\n",
        "  img = PIL.Image.open(io.BytesIO(r.content))\n",
        "  img.thumbnail((max_size, max_size), PIL.Image.ANTIALIAS)\n",
        "  img = np.float32(img)/255.0\n",
        "  # premultiply RGB by Alpha\n",
        "  img[..., :3] *= img[..., 3:]\n",
        "  return img\n",
        "\n",
        "def to_rgba(x):\n",
        "  return x[..., :4]\n",
        "\n",
        "def to_alpha(x):\n",
        "  return tf.clip_by_value(x[..., 3:4], 0.0, 1.0)\n",
        "\n",
        "def to_rgb(x):\n",
        "  # assume rgb premultiplied by alpha\n",
        "  rgb, a = x[..., :3], to_alpha(x)\n",
        "  return 1.0-a+rgb\n",
        "\n",
        "def get_living_mask(x):\n",
        "  alpha = x[:, :, :, 3:4]\n",
        "  return tf.nn.max_pool2d(alpha, 3, [1, 1, 1, 1], 'SAME') \u003e 0.1\n",
        "\n",
        "def make_seed(size, n=1):\n",
        "  x = np.zeros([n, size, size, CHANNEL_N], np.float32)\n",
        "  x[:, size//2, size//2, 3:] = 1.0\n",
        "  return x\n",
        "\n",
        "\n",
        "class CAModel(tf.keras.Model):\n",
        "\n",
        "  def __init__(self, channel_n=CHANNEL_N, fire_rate=CELL_FIRE_RATE):\n",
        "    super().__init__()\n",
        "    self.channel_n = channel_n\n",
        "    self.fire_rate = fire_rate\n",
        "\n",
        "    self.dmodel = tf.keras.Sequential([\n",
        "          Conv2D(128, 1, activation=tf.nn.relu),\n",
        "          Conv2D(self.channel_n, 1, activation=None,\n",
        "              kernel_initializer=tf.zeros_initializer),\n",
        "    ])\n",
        "\n",
        "    self(tf.zeros([1, 3, 3, channel_n]))  # dummy call to build the model\n",
        "\n",
        "  @tf.function\n",
        "  def perceive(self, x, angle=0.0):\n",
        "    identify = np.float32([0, 1, 0])\n",
        "    identify = np.outer(identify, identify)\n",
        "    dx = np.outer([1, 2, 1], [-1, 0, 1]) / 8.0  # Sobel filter\n",
        "    dy = dx.T\n",
        "    c, s = tf.cos(angle), tf.sin(angle)\n",
        "    kernel = tf.stack([identify, c*dx-s*dy, s*dx+c*dy], -1)[:, :, None, :]\n",
        "    kernel = tf.repeat(kernel, self.channel_n, 2)\n",
        "    y = tf.nn.depthwise_conv2d(x, kernel, [1, 1, 1, 1], 'SAME')\n",
        "    return y\n",
        "\n",
        "  @tf.function\n",
        "  def call(self, x, fire_rate=None, angle=0.0, step_size=1.0):\n",
        "    pre_life_mask = get_living_mask(x)\n",
        "\n",
        "    y = self.perceive(x, angle)\n",
        "    dx = self.dmodel(y)*step_size\n",
        "    if fire_rate is None:\n",
        "      fire_rate = self.fire_rate\n",
        "    update_mask = tf.random.uniform(tf.shape(x[:, :, :, :1])) \u003c= fire_rate\n",
        "    x += dx * tf.cast(update_mask, tf.float32)\n",
        "\n",
        "    post_life_mask = get_living_mask(x)\n",
        "    life_mask = pre_life_mask \u0026 post_life_mask\n",
        "    return x * tf.cast(life_mask, tf.float32)\n",
        "\n",
        "\n",
        "CAModel().dmodel.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IeWf6HeTe8kM"
      },
      "outputs": [],
      "source": [
        "#@title Train Utilities (SamplePool, Model Export, Damage)\n",
        "from google.protobuf.json_format import MessageToDict\n",
        "from tensorflow.python.framework import convert_to_constants\n",
        "\n",
        "class SamplePool:\n",
        "  def __init__(self, *, _parent=None, _parent_idx=None, **slots):\n",
        "    self._parent = _parent\n",
        "    self._parent_idx = _parent_idx\n",
        "    self._slot_names = slots.keys()\n",
        "    self._size = None\n",
        "    for k, v in slots.items():\n",
        "      if self._size is None:\n",
        "        self._size = len(v)\n",
        "      assert self._size == len(v)\n",
        "      setattr(self, k, np.asarray(v))\n",
        "\n",
        "  def sample(self, n):\n",
        "    idx = np.random.choice(self._size, n, False)\n",
        "    batch = {k: getattr(self, k)[idx] for k in self._slot_names}\n",
        "    batch = SamplePool(**batch, _parent=self, _parent_idx=idx)\n",
        "    return batch\n",
        "\n",
        "  def commit(self):\n",
        "    for k in self._slot_names:\n",
        "      getattr(self._parent, k)[self._parent_idx] = getattr(self, k)\n",
        "\n",
        "@tf.function\n",
        "def make_circle_masks(n, h, w):\n",
        "  x = tf.linspace(-1.0, 1.0, w)[None, None, :]\n",
        "  y = tf.linspace(-1.0, 1.0, h)[None, :, None]\n",
        "  center = tf.random.uniform([2, n, 1, 1], -0.5, 0.5)\n",
        "  r = tf.random.uniform([n, 1, 1], 0.1, 0.4)\n",
        "  x, y = (x-center[0])/r, (y-center[1])/r\n",
        "  mask = tf.cast(x*x+y*y \u003c 1.0, tf.float32)\n",
        "  return mask\n",
        "\n",
        "def export_model(ca, base_fn):\n",
        "  ca.save_weights(base_fn)\n",
        "\n",
        "  cf = ca.call.get_concrete_function(\n",
        "      x=tf.TensorSpec([None, None, None, CHANNEL_N]),\n",
        "      fire_rate=tf.constant(0.5),\n",
        "      angle=tf.constant(0.0),\n",
        "      step_size=tf.constant(1.0))\n",
        "  cf = convert_to_constants.convert_variables_to_constants_v2(cf)\n",
        "  graph_def = cf.graph.as_graph_def()\n",
        "  graph_json = MessageToDict(graph_def)\n",
        "  graph_json['versions'] = dict(producer='1.14', minConsumer='1.14')\n",
        "  model_json = {\n",
        "      'format': 'graph-model',\n",
        "      'modelTopology': graph_json,\n",
        "      'weightsManifest': [],\n",
        "  }\n",
        "  with open(base_fn+'.json', 'w') as f:\n",
        "    json.dump(model_json, f)\n",
        "\n",
        "def generate_pool_figures(pool, step_i):\n",
        "  tiled_pool = tile2d(to_rgb(pool.x[:49]))\n",
        "  fade = np.linspace(1.0, 0.0, 72)\n",
        "  ones = np.ones(72) \n",
        "  tiled_pool[:, :72] += (-tiled_pool[:, :72] + ones[None, :, None]) * fade[None, :, None] \n",
        "  tiled_pool[:, -72:] += (-tiled_pool[:, -72:] + ones[None, :, None]) * fade[None, ::-1, None]\n",
        "  tiled_pool[:72, :] += (-tiled_pool[:72, :] + ones[:, None, None]) * fade[:, None, None]\n",
        "  tiled_pool[-72:, :] += (-tiled_pool[-72:, :] + ones[:, None, None]) * fade[::-1, None, None]\n",
        "  imwrite('train_log/%04d_pool.jpg'%step_i, tiled_pool)\n",
        "\n",
        "def visualize_batch(x0, x, step_i):\n",
        "  vis0 = np.hstack(to_rgb(x0).numpy())\n",
        "  vis1 = np.hstack(to_rgb(x).numpy())\n",
        "  vis = np.vstack([vis0, vis1])\n",
        "  imwrite('train_log/batches_%04d.jpg'%step_i, vis)\n",
        "  print('batch (before/after):')\n",
        "  imshow(vis)\n",
        "\n",
        "def plot_loss(loss_log):\n",
        "  pl.figure(figsize=(10, 4))\n",
        "  pl.title('Loss history (log10)')\n",
        "  pl.plot(np.log10(loss_log), '.', alpha=0.1)\n",
        "  pl.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8VWuk3GqxKS7"
      },
      "outputs": [],
      "source": [
        "!wget -O models.zip 'https://github.com/google-research/self-organising-systems/blob/master/assets/growing_ca/models.zip?raw=true'\n",
        "!unzip -oq models.zip\n",
        "\n",
        "# available pretrained emoji are:\n",
        "#EMOJI = '🦎😀💥👁🐠🦋🐞🕸🥨🎄'\n",
        "\n",
        "def get_model(emoji='🦎', fire_rate=0.5, use_pool=1, damage_n=3, run=0,\n",
        "              prefix='models/', output='model'):\n",
        "  path = prefix\n",
        "  assert fire_rate in [0.5, 1.0]\n",
        "  if fire_rate==0.5:\n",
        "    path += 'use_sample_pool_%d damage_n_%d '%(use_pool, damage_n)\n",
        "  elif fire_rate==1.0:\n",
        "    path += 'fire_rate_1.0 '\n",
        "  code = hex(ord(emoji))[2:].upper()\n",
        "  path += 'target_emoji_%s run_index_%d/08000'%(code, run)\n",
        "  assert output in ['model', 'json']\n",
        "  if output == 'model':\n",
        "    ca = CAModel(channel_n=16, fire_rate=fire_rate)\n",
        "    ca.load_weights(path)\n",
        "    return ca\n",
        "  elif output == 'json':\n",
        "    return open(path+'.json', 'r').read()\n",
        "\n",
        "def get_local_model(path, output='model'):\n",
        "  assert output in ['model', 'json']\n",
        "  if output == 'model':\n",
        "    ca = CAModel(channel_n=16)\n",
        "    ca.load_weights(path)\n",
        "    return ca\n",
        "  elif output == 'json':\n",
        "    return open(path+'.json', 'r').read()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nFy_S-scxiDP"
      },
      "source": [
        "## CA Targets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PkI-1mTHGM1I"
      },
      "outputs": [],
      "source": [
        "!wget -O growing_ca_target_images.zip 'https://github.com/google-research/self-organising-systems/blob/master/adversarial_reprogramming_ca/assets/growing_ca_target_images.zip?raw=true'\n",
        "!unzip -oq \"growing_ca_target_images.zip\" -d \"target_pics\"\n",
        "!ls target_pics/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KXyeAt2zyUfW"
      },
      "outputs": [],
      "source": [
        "# If you want a new picture, modify it with either:\n",
        "# - https://pixlr.com/\n",
        "# - https://www.piskelapp.com/\n",
        "def load_image_from_file(fp, max_size=TARGET_SIZE):\n",
        "  img = PIL.Image.open(fp)\n",
        "  img.thumbnail((max_size, max_size), PIL.Image.ANTIALIAS)\n",
        "  img = np.float32(img)/255.0\n",
        "  # premultiply RGB by Alpha\n",
        "  img[..., :3] *= img[..., 3:]\n",
        "  return img\n",
        "\n",
        "pic_prefix = \"target_pics/\"\n",
        "lizard_no_tail_fp = pic_prefix + \"lizard_no_tail.png\"\n",
        "lizard_no_leg_fp = pic_prefix + \"lizard_no_leg.png\"\n",
        "lizard_no_head_fp = pic_prefix + \"lizard_no_head.png\"\n",
        "lizard_no_arm_fp = pic_prefix + \"lizard_no_arm.png\"\n",
        "lizard_red_fp = pic_prefix + \"lizard_red.png\"\n",
        "lizard_blue_fp = pic_prefix + \"lizard_blue.png\"\n",
        "target_map = {\n",
        "    \"liz_no_tail\": lizard_no_tail_fp,\n",
        "    \"lizard_no_leg\": lizard_no_leg_fp,\n",
        "    \"lizard_no_head\": lizard_no_head_fp,\n",
        "    \"lizard_no_arm\": lizard_no_arm_fp,\n",
        "    \"lizard_red\": lizard_red_fp,\n",
        "    \"lizard_blue\": lizard_blue_fp}\n",
        "\n",
        "lizard_complete_fp = pic_prefix + \"lizard_complete.png\"\n",
        "\n",
        "#butterfly_all_orange_fp = pic_prefix + \"butterfly_all_orange.png\"\n",
        "#butterfly_left_orange_fp = pic_prefix + \"butterfly_left_orange.png\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ByHbsY0EuyqB"
      },
      "source": [
        "# Experiment: Virus takeover\n",
        "\n",
        "We get a pretrained model (lizard) and insert a small number of cells running a different CA rule. This CA needs to reprogram the entire CA group to create a different lizard."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MImpr11oqamb"
      },
      "outputs": [],
      "source": [
        "target_label = 'liz_no_tail' #@param [\"liz_no_tail\", \"liz_no_leg\", \"liz_no_head\", \"lizard_no_arm\", \"lizard_red\", \"lizard_blue\"]\n",
        "\n",
        "target_img = load_image_from_file(target_map[target_label])\n",
        "imshow(zoom(to_rgb(target_img), 2), fmt='png')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pxvsFHTNhTXm"
      },
      "outputs": [],
      "source": [
        "def generate_padded_xy(p, target_img, ca):\n",
        "  pad_target = tf.pad(target_img, [(p, p), (p, p), (0, 0)])\n",
        "  h, w = pad_target.shape[:2]\n",
        "  seed = np.zeros([h, w, CHANNEL_N], np.float32)\n",
        "  seed[h//2, w//2, 3:] = 1.0\n",
        "  x0_seed = seed[None, ...]\n",
        "  for i in tf.range(200):\n",
        "    x0_seed = ca(x0_seed)\n",
        "  return x0_seed, pad_target"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i7OSde31xbGK"
      },
      "outputs": [],
      "source": [
        "#@title Initialize Training { vertical-output: true}\n",
        "\n",
        "#@markdown Recommended: 0.1 percentage_virus for no_tail, and 0.6 for red lizard.\n",
        "percentage_virus = 0.1 #@param [0.1, 0.6] {allow-input: true}\n",
        "\n",
        "# Number of pixels used to pad the target image border.\n",
        "# Use 32 for the no_tail experiment\n",
        "TARGET_PADDING = 32\n",
        "# YOU CAN use 16 for the red experiment\n",
        "# TARGET_PADDING = 16\n",
        "\n",
        "def loss_f(x):\n",
        "  return tf.reduce_mean(tf.square(to_rgba(x)-pad_target), [-2, -3, -1])\n",
        "\n",
        "target_emoji = '🦎' # '🦋' \n",
        "ca = get_model(emoji=target_emoji)\n",
        "x0_seed, pad_target = generate_padded_xy(TARGET_PADDING, target_img, ca)\n",
        "h, w = pad_target.shape[:2]\n",
        "\n",
        "print(\"This is the initial configuration\")\n",
        "imshow(zoom(to_rgb(x0_seed[0]), 2), fmt='png')\n",
        "\n",
        "\n",
        "# We create a mask: you can't place viruses outside this.\n",
        "print(\"Mask visualization.\")\n",
        "pad_target_mask = tf.cast(pad_target[...,3:4] \u003e= 0.1, tf.float32)\n",
        "pl.imshow(pad_target_mask[:,:,0])\n",
        "pl.show()\n",
        "\n",
        "@tf.function\n",
        "def random_mask_not_out_target(bsize, perc, pad_target_mask):\n",
        "  h, w = pad_target_mask.shape[0:2]\n",
        "  mask = tf.cast(\n",
        "    tf.random.uniform([bsize, h, w, 1]) \u003c perc, \n",
        "    tf.float32)\n",
        "  return mask * pad_target_mask\n",
        "\n",
        "print(\"Example target with random mask.\")\n",
        "random_mask = random_mask_not_out_target(1, percentage_virus, pad_target_mask)\n",
        "imshow(zoom(to_rgb(x0_seed[0] * random_mask[0]), 2), fmt='png')\n",
        "\n",
        "adv_mask = random_mask\n",
        "inv_adv_mask = 1. - adv_mask\n",
        "\n",
        "loss_log = []\n",
        "\n",
        "lr = 2e-3\n",
        "lr_sched = tf.keras.optimizers.schedules.PiecewiseConstantDecay(\n",
        "    [2000], [lr, lr*0.1])\n",
        "trainer = tf.keras.optimizers.Adam(lr_sched)\n",
        "\n",
        "loss0 = loss_f(x0_seed).numpy()\n",
        "\n",
        "initial_pool_elems = np.repeat(x0_seed, POOL_SIZE, 0)\n",
        "initial_random_mask = random_mask_not_out_target(\n",
        "    POOL_SIZE, percentage_virus, pad_target_mask).numpy()\n",
        "pool = SamplePool(x=initial_pool_elems, m=initial_random_mask)\n",
        "\n",
        "# The spooky virus\n",
        "adversarial_ca = CAModel()\n",
        "\n",
        "!mkdir -p train_log \u0026\u0026 rm -f train_log/*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ae9OY1gK4JZ9"
      },
      "outputs": [],
      "source": [
        "#@title Training Loop {vertical-output: true}\n",
        "\n",
        "@tf.function\n",
        "def train_step(x, m):\n",
        "  iter_n = tf.random.uniform([], 64, 96, tf.int32)\n",
        "  with tf.GradientTape() as g:\n",
        "    for i in tf.range(iter_n):\n",
        "      x_orig = ca(x)\n",
        "      x_adv = adversarial_ca(x)\n",
        "      x = x_orig * (1. - m) + x_adv * m\n",
        "    loss = tf.reduce_mean(loss_f(x))\n",
        "  grads = g.gradient(loss, adversarial_ca.weights)\n",
        "  grads = [g/(tf.norm(g)+1e-8) for g in grads]\n",
        "  trainer.apply_gradients(zip(grads, adversarial_ca.weights))\n",
        "  return x, loss\n",
        "\n",
        "for i in range(100000+1):\n",
        "  batch = pool.sample(BATCH_SIZE)\n",
        "  x0 = batch.x\n",
        "  m0 = batch.m\n",
        "  loss_rank = loss_f(x0).numpy().argsort()[::-1]\n",
        "  x0 = x0[loss_rank]\n",
        "  m0 = m0[loss_rank]\n",
        "  x0[:1] = x0_seed[0]\n",
        "  m0[:1] = random_mask_not_out_target(\n",
        "      1, percentage_virus, pad_target_mask)[0].numpy()\n",
        "\n",
        "  if DAMAGE_N:\n",
        "    damage = 1.0-make_circle_masks(DAMAGE_N, h, w).numpy()[..., None]\n",
        "    x0[-DAMAGE_N:] *= damage\n",
        "\n",
        "  x, loss = train_step(x0, m0)\n",
        "\n",
        "  batch.x[:] = x\n",
        "  batch.m[:] = m0\n",
        "  batch.commit()\n",
        "\n",
        "  step_i = len(loss_log)\n",
        "  loss_log.append(loss.numpy())\n",
        "  \n",
        "  if step_i%10 == 0:\n",
        "    generate_pool_figures(pool, step_i)\n",
        "  if step_i%100 == 0:\n",
        "    clear_output()\n",
        "    visualize_batch(x0, x, step_i)\n",
        "    plot_loss(loss_log)\n",
        "  if step_i%1000 == 0:\n",
        "    export_model(adversarial_ca, 'train_log/%06d'%step_i)\n",
        "\n",
        "  print('\\r step: %d, log10(loss): %.3f'%(len(loss_log), np.log10(loss)), end='')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BnzV4k3Uw5Zc"
      },
      "outputs": [],
      "source": [
        "# useful code if you end up interrupting the run.\n",
        "print(step_i)\n",
        "export_model(adversarial_ca, 'train_log/%06d'%step_i)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5ZeYEvJ2Hxv1"
      },
      "outputs": [],
      "source": [
        "# @title Load saved models from Github\n",
        "!wget -O growing_ca_adversarial_models.zip 'https://github.com/google-research/self-organising-systems/blob/master/adversarial_reprogramming_ca/assets/growing_ca_adversarial_models.zip?raw=true'\n",
        "!unzip -oq \"growing_ca_adversarial_models.zip\" -d \"saved_adversarial\"\n",
        "\n",
        "no_tail_ca = get_local_model(\"saved_adversarial/no_tail_model\", output='model')\n",
        "red_10p_ca = get_local_model(\"saved_adversarial/red_model\", output='model')\n",
        "red_60p_ca = get_local_model(\"saved_adversarial/red_model_60perc\", output='model')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SGIWYWiqxRN6"
      },
      "outputs": [],
      "source": [
        "#@title TensorFlow.js Demo {run:\"auto\", vertical-output: true}\n",
        "#@markdown Select \"CHECKPOINT\" model to load the checkpoint created by running cells from the \"Training\" section of this notebook\n",
        "\n",
        "#@markdown Shift-click to seed the pattern\n",
        "import IPython.display\n",
        "\n",
        "# These are the original models' parameters. Modify them manually\n",
        "# if you want to see different behaviors.\n",
        "model = \"\\uD83E\\uDD8E 1F98E\"  #[😀 1F600', '💥 1F4A5', '👁 1F441', '🦎 1F98E', '🐠 1F420', '🦋 1F98B', '🐞 1F41E', '🕸 1F578', '🥨 1F968', '🎄 1F384']\n",
        "model_type = '3 regenerating'  #['1 naive', '2 persistent', '3 regenerating']\n",
        "\n",
        "code = model.split(' ')[1]\n",
        "emoji = chr(int(code, 16))\n",
        "experiment_i = int(model_type.split()[0])-1\n",
        "use_pool = (0, 1, 1)[experiment_i]\n",
        "damage_n = (0, 0, 3)[experiment_i]\n",
        "model_str = get_model(emoji, use_pool=use_pool, damage_n=damage_n, output='json')\n",
        "\n",
        "adversarial_model_source = \"no_tail_model\"  #@param ['CHECKPOINT', 'no_tail_model', 'red_model', 'red_model_60perc' ]\n",
        "\n",
        "if adversarial_model_source != \"CHECKPOINT\":\n",
        "  adv_model_str = get_local_model(\"saved_adversarial/\" + adversarial_model_source,output=\"json\")\n",
        "else:\n",
        "  last_checkpoint_fn = sorted(glob.glob('train_log/*.json'))[-1]\n",
        "  adv_model_str = open(last_checkpoint_fn).read()\n",
        "\n",
        "\n",
        "data_js = '''\n",
        "  window.GRAPH_URL = URL.createObjectURL(new Blob([`%s`], {type: 'application/json'}));\n",
        "  window.ADV_GRAPH_URL = URL.createObjectURL(new Blob([`%s`], {type: 'application/json'}));\n",
        "'''%(model_str, adv_model_str)\n",
        "\n",
        "display(IPython.display.Javascript(data_js))\n",
        "\n",
        "\n",
        "IPython.display.HTML('''\n",
        "\u003cscript src=\"https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@2.3.0/dist/tf.min.js\"\u003e\u003c/script\u003e\n",
        "\u003cscript src=\"https://cdnjs.cloudflare.com/ajax/libs/cash/4.1.2/cash.min.js\"\u003e\u003c/script\u003e\n",
        "\n",
        "\u003ccanvas id='canvas' style=\"border: 1px solid black; image-rendering: pixelated;\"\u003e\u003c/canvas\u003e\n",
        "\u003cdiv\u003e\u003cbutton type=\"button\" id=\"removeadv\"\u003eRemove adversaries\u003c/button\u003e\u003c/div\u003e\n",
        "\n",
        "\u003cdiv class=\"boxcontainer\"\u003e\n",
        "\u003cinput type=\"checkbox\" id=\"drawadversary\" name=\"drawadversary\"\u003e\n",
        "\u003clabel for=\"drawadversary\"\u003eDraw adversary\u003c/label\u003e\u003cbr\u003e\n",
        "\u003c/div\u003e\n",
        "\u003cdiv class=\"boxcontainer\"\u003e\n",
        "\u003cinput type=\"checkbox\" id=\"showadvmask\" name=\"showadvmask\"\u003e\n",
        "\u003clabel for=\"showadvmask\"\u003eShow adversary mask\u003c/label\u003e\u003cbr\u003e\n",
        "\u003c/div\u003e\n",
        "\n",
        "\u003cscript\u003e\n",
        "  \"use strict\";\n",
        "  \n",
        "  const sleep = (ms)=\u003enew Promise(resolve =\u003e setTimeout(resolve, ms));\n",
        "  \n",
        "  const parseConsts = model_graph=\u003e{\n",
        "    const dtypes = {'DT_INT32':['int32', 'intVal', Int32Array],\n",
        "                    'DT_FLOAT':['float32', 'floatVal', Float32Array]};\n",
        "    \n",
        "    const consts = {};\n",
        "    model_graph.modelTopology.node.filter(n=\u003en.op=='Const').forEach((node=\u003e{\n",
        "      const v = node.attr.value.tensor;\n",
        "      const [dtype, field, arrayType] = dtypes[v.dtype];\n",
        "      if (!v.tensorShape.dim) {\n",
        "        consts[node.name] = [tf.scalar(v[field][0], dtype)];\n",
        "      } else {\n",
        "        // if there is a 0-length dimension, the exported graph json lacks \"size\"\n",
        "        const shape = v.tensorShape.dim.map(d=\u003e(!d.size) ? 0 :parseInt(d.size));\n",
        "        let arr;\n",
        "        if (v.tensorContent) {\n",
        "          const data = atob(v.tensorContent);\n",
        "          const buf = new Uint8Array(data.length);\n",
        "          for (var i=0; i\u003cdata.length; ++i) {\n",
        "            buf[i] = data.charCodeAt(i);\n",
        "          }\n",
        "          arr = new arrayType(buf.buffer);\n",
        "        } else {\n",
        "          const size = shape.reduce((a, b)=\u003ea*b);\n",
        "          arr = new arrayType(size);\n",
        "          if (size){\n",
        "            arr.fill(v[field][0]);\n",
        "          }\n",
        "        }\n",
        "        consts[node.name] = [tf.tensor(arr, shape, dtype)];\n",
        "      }\n",
        "    }));\n",
        "    return consts;\n",
        "  }\n",
        "  \n",
        "  let drawadversaryCkbx = document.getElementById(\"drawadversary\");\n",
        "  let showadvmaskCkbx = document.getElementById(\"showadvmask\");\n",
        "\n",
        "  const run = async ()=\u003e{\n",
        "    const r = await fetch(GRAPH_URL);\n",
        "    const consts = parseConsts(await r.json());\n",
        "    \n",
        "    const model = await tf.loadGraphModel(GRAPH_URL);\n",
        "    Object.assign(model.weights, consts);\n",
        "    \n",
        "    // Adversarial model now.\n",
        "    const rad = await fetch(ADV_GRAPH_URL);\n",
        "    const adv_consts = parseConsts(await rad.json());\n",
        "    const adv_model = await tf.loadGraphModel(ADV_GRAPH_URL);\n",
        "\n",
        "    console.log(\"Loaded adv model\")\n",
        "    Object.assign(adv_model.weights, adv_consts);\n",
        "    \n",
        "    let seed = new Array(16).fill(0).map((x, i)=\u003ei\u003c3?0:1);\n",
        "    seed = tf.tensor(seed, [1, 1, 1, 16]);\n",
        "    \n",
        "    const D = 96;\n",
        "    const initState = tf.tidy(()=\u003e{\n",
        "      const D2 = D/2;\n",
        "      const a = seed.pad([[0, 0], [D2-1, D2], [D2-1, D2], [0,0]]);\n",
        "      return a;\n",
        "    });\n",
        "    \n",
        "    const state = tf.variable(initState);\n",
        "    // this is where we keep track of where is which CA.\n",
        "    const adv_mask = tf.variable(tf.zeros([1, D, D, 1]));\n",
        "    // store this to avoid recomputations.\n",
        "    const orig_mask = tf.variable(tf.ones([1, D, D, 1]));\n",
        "    const [_, h, w, ch] = state.shape;\n",
        "    \n",
        "    const damage = (x, y, r)=\u003e{\n",
        "      tf.tidy(()=\u003e{\n",
        "        const rx = tf.range(0, w).sub(x).div(r).square().expandDims(0);\n",
        "        const ry = tf.range(0, h).sub(y).div(r).square().expandDims(1);\n",
        "        const mask = rx.add(ry).greater(1.0).expandDims(2);\n",
        "        state.assign(state.mul(mask));\n",
        "      });\n",
        "    }\n",
        "    \n",
        "    const plantSeed = (x, y)=\u003e{\n",
        "      const x2 = w-x-seed.shape[2];\n",
        "      const y2 = h-y-seed.shape[1];\n",
        "      if (x\u003c0 || x2\u003c0 || y2\u003c0 || y2\u003c0)\n",
        "        return;\n",
        "      tf.tidy(()=\u003e{\n",
        "        const a = seed.pad([[0, 0], [y, y2], [x, x2], [0,0]]);\n",
        "        state.assign(state.add(a));\n",
        "      });\n",
        "    }\n",
        "\n",
        "    $('#removeadv').on('click', e=\u003e{\n",
        "        tf.tidy(()=\u003e{\n",
        "          adv_mask.assign(tf.zeros([1, D, D, 1]));\n",
        "          orig_mask.assign(tf.ones([1, D, D, 1]));\n",
        "        });\n",
        "    });\n",
        "\n",
        "    \n",
        "    const scale = 4;\n",
        "    \n",
        "    const canvas = document.getElementById('canvas');\n",
        "    const ctx = canvas.getContext('2d');\n",
        "    canvas.width = w;\n",
        "    canvas.height = h;\n",
        "    canvas.style.width = `${w*scale}px`;\n",
        "    canvas.style.height = `${h*scale}px`;\n",
        "\n",
        "    const getClickPos = e=\u003e{\n",
        "        const x = Math.floor((e.pageX-e.target.offsetLeft) / scale);\n",
        "        const y = Math.floor((e.pageY-e.target.offsetTop) / scale);\n",
        "        return [x, y];\n",
        "    }\n",
        "\n",
        "    canvas.onmousedown = e=\u003e{\n",
        "      const [x, y] = getClickPos(e);\n",
        "      //const x = Math.floor(e.clientX/scale);\n",
        "      //const y = Math.floor(e.clientY/scale);\n",
        "      if (drawadversaryCkbx.checked) {\n",
        "        // perform surgical insertions!\n",
        "        tf.tidy(()=\u003e {\n",
        "          const insertion = tf.ones([1, 1]).pad(\n",
        "            [[y, h-y-1], [x, w-x-1]]).expandDims(0).expandDims(3);\n",
        "          const mask = tf.tensor(1.).sub(insertion);\n",
        "          adv_mask.assign(adv_mask.mul(mask).add(insertion));\n",
        "          orig_mask.assign(adv_mask.sub(1).mul(-1));\n",
        "        });\n",
        "        return;\n",
        "      }\n",
        "      if (e.buttons == 1) {\n",
        "        if (e.shiftKey) {\n",
        "          plantSeed(x, y);  \n",
        "        } else {\n",
        "          damage(x, y, 8);\n",
        "        }\n",
        "      }\n",
        "    }\n",
        "    canvas.onmousemove = e=\u003e{\n",
        "      const x = Math.floor(e.clientX/scale);\n",
        "      const y = Math.floor(e.clientY/scale);\n",
        "      if (e.buttons == 1 \u0026\u0026 !e.shiftKey \u0026\u0026 !drawadversaryCkbx.checked) {\n",
        "        damage(x, y, 8);\n",
        "      }\n",
        "    }\n",
        "\n",
        "    function step() {\n",
        "      tf.tidy(()=\u003e{\n",
        "        const orig_state = model.execute(\n",
        "            {x:state, fire_rate:tf.tensor(0.5),\n",
        "            angle:tf.tensor(0.0), step_size:tf.tensor(1.0)}, \n",
        "            ['Identity']);\n",
        "        const adv_state = adv_model.execute(\n",
        "            {x:state, fire_rate:tf.tensor(0.5),\n",
        "            angle:tf.tensor(0.0), step_size:tf.tensor(1.0)}, \n",
        "            ['Identity']);\n",
        "\n",
        "        state.assign(orig_state.mul(orig_mask).add(adv_state.mul(adv_mask)));\n",
        "      });\n",
        "    }\n",
        "\n",
        "    const initT = new Date().getTime() / 1000;\n",
        "    function render() {\n",
        "      step();\n",
        "\n",
        "      const imageData = tf.tidy(()=\u003e{\n",
        "        const rgba = state.slice([0, 0, 0, 0], [-1, -1, -1, 4]);\n",
        "        const a = state.slice([0, 0, 0, 3], [-1, -1, -1, 1]);\n",
        "        let img = tf.tensor(1.0).sub(a).add(rgba).mul(255);\n",
        "\n",
        "        const advred = adv_mask.gather([0], 3) // R\n",
        "                  .pad([[0, 0], [0, 0], [0, 0], [0, 2]], 0) // GB\n",
        "                  .pad([[0, 0], [0, 0], [0, 0], [0, 1]], 1) // A\n",
        "                  .mul(255);\n",
        "        const seconds = new Date().getTime() / 1000 - initT;\n",
        "        const t = tf.tensor(seconds).sin().abs();\n",
        "        const onemt = tf.tensor(1.).sub(t);\n",
        "        const adv_mask_period = adv_mask.mul(t);\n",
        "        const img_behind_mask = img.mul(adv_mask);\n",
        "        const adv_image = advred.mul(adv_mask_period).add(\n",
        "          img_behind_mask.mul(onemt));\n",
        "        //const orig_mask_period = orig_mask.mul(tf.tensor(seconds).sin().abs());\n",
        "        //const adv_mask_period = tf.tensor(1.0).sub(orig_mask_period)\n",
        "        img = img.mul(orig_mask).add(adv_image);\n",
        "\n",
        "        const rgbaBytes = new Uint8ClampedArray(img.dataSync());\n",
        "        return new ImageData(rgbaBytes, w, h);\n",
        "      });\n",
        "      ctx.putImageData(imageData, 0, 0);\n",
        "\n",
        "      requestAnimationFrame(render);\n",
        "    }\n",
        "    render();\n",
        "  }\n",
        "  run();\n",
        "  \n",
        "\u003c/script\u003e\n",
        "''')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z20YawBrI0Qo"
      },
      "source": [
        "## Figures"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M6LPCoA0Tihj"
      },
      "outputs": [],
      "source": [
        "#@title viz targets\n",
        "liz_compl = load_image_from_file(lizard_complete_fp)\n",
        "imshow(zoom(to_rgb(liz_compl), 2), fmt='png')\n",
        "\n",
        "liz_notail = load_image_from_file(lizard_no_tail_fp)\n",
        "liz_red = load_image_from_file(lizard_red_fp)\n",
        "imshow(zoom(to_rgb(np.hstack([liz_notail, liz_red])), 2), fmt='png')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J9rA_Zms81oI"
      },
      "outputs": [],
      "source": [
        "#@title viz current run\n",
        "import PIL.ImageFont\n",
        "\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "x = x0_seed\n",
        "exp_percentage_virus = 0.1\n",
        "new_mask = random_mask_not_out_target(1, exp_percentage_virus).numpy()\n",
        "inv_new_mask = 1. - new_mask\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(10001):\n",
        "    if i%20 == 0:\n",
        "      # We want to show something different on step 500:\n",
        "      if i == 500 or i == 1000:\n",
        "        x_curr = x[0] * inv_new_mask[0]\n",
        "      else:\n",
        "        x_curr = x[0]\n",
        "      vis = zoom(to_rgb(x_curr), 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perc virus: {}%\".format(int(exp_percentage_virus * 100))\n",
        "      perturbcolor = (255,0,0)# if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "      vid.add(np.uint8(im))\n",
        "      if i == 500 or i == 1000:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(50):\n",
        "          vid.add(np.uint8(im))\n",
        "    x_orig = orig_ca(x)\n",
        "    x_adv = adversarial_ca(x)\n",
        "    x = x_orig * inv_new_mask + x_adv * new_mask\n",
        "\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hIAeOVQx3r5O"
      },
      "outputs": [],
      "source": [
        "#@title viz no tail saved run\n",
        "import PIL.ImageFont\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "orig_ca = get_model(target_emoji)\n",
        "\n",
        "no_tail_target_img = load_image_from_file(lizard_no_tail_fp)\n",
        "p = 24\n",
        "exp_x0_seed, no_tail_pad_target = generate_padded_xy(p, no_tail_target_img, orig_ca)\n",
        "\n",
        "no_tail_pad_target_mask = tf.cast(no_tail_pad_target[...,3:4] \u003e= 0.1, tf.float32)\n",
        "\n",
        "exp_percentage_virus = 0.1\n",
        "num_runs = 6\n",
        "x =  np.repeat(exp_x0_seed, num_runs, 0)\n",
        "wtile = 3\n",
        "new_mask = random_mask_not_out_target(num_runs, exp_percentage_virus, \n",
        "                                      no_tail_pad_target_mask).numpy()\n",
        "inv_new_mask = 1. - new_mask\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(10001):\n",
        "    if i%20 == 0:\n",
        "      \"\"\"\n",
        "      # We want to show something different on step 500:\n",
        "      if i == 500 or i == 1000:\n",
        "        x_curr = x[0] * inv_new_mask[0]\n",
        "      else:\n",
        "        x_curr = x[0]\n",
        "      \"\"\"\n",
        "      x_curr = x\n",
        "      rgb = to_rgb(x_curr)\n",
        "      # Make it blinking with adversaries!\n",
        "      adv_rgb_color = tf.constant([[[[1.,0., 0.]]]])\n",
        "      t = 0.5 + tf.sin(i/20 / 10) / 2.\n",
        "      adv_mask_period = new_mask * t\n",
        "      img_behind_mask = rgb * new_mask\n",
        "      adv_rgb = adv_rgb_color * adv_mask_period + img_behind_mask * (1 - t)\n",
        "\n",
        "      rgb = rgb * inv_new_mask + adv_rgb\n",
        "\n",
        "      if num_runs == 1:\n",
        "        rgb = rgb[0]\n",
        "      else:\n",
        "        rgb = tile2d(rgb, w=wtile)\n",
        "\n",
        "\n",
        "\n",
        "      vis = zoom(rgb, 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perc virus: {}%\".format(int(exp_percentage_virus * 100))\n",
        "      perturbcolor = (255,0,0)# if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "      vid.add(np.uint8(im))\n",
        "      \"\"\"\n",
        "      if i == 500 or i == 1000:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(50):\n",
        "          vid.add(np.uint8(im))\n",
        "      \"\"\"\n",
        "    x_orig = orig_ca(x)\n",
        "    x_adv = no_tail_ca(x)\n",
        "    x = x_orig * inv_new_mask + x_adv * new_mask\n",
        "\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E12pjhDd-4y4"
      },
      "outputs": [],
      "source": [
        "#@title viz red 10% saved run\n",
        "import PIL.ImageFont\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "orig_ca = get_model(target_emoji)\n",
        "adv_ca = red_10p_ca\n",
        "\n",
        "red_target_img = load_image_from_file(lizard_red_fp)\n",
        "p = 24\n",
        "exp_x0_seed, red_pad_target = generate_padded_xy(p, red_target_img, orig_ca)\n",
        "\n",
        "red_pad_target_mask = tf.cast(red_pad_target[...,3:4] \u003e= 0.1, tf.float32)\n",
        "\n",
        "exp_percentage_virus = 0.1\n",
        "num_runs = 6\n",
        "x =  np.repeat(exp_x0_seed, num_runs, 0)\n",
        "wtile = 3\n",
        "new_mask = random_mask_not_out_target(num_runs, exp_percentage_virus, \n",
        "                                      red_pad_target_mask).numpy()\n",
        "inv_new_mask = 1. - new_mask\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(1001):\n",
        "    if i%5 == 0:\n",
        "      \n",
        "      # We want to show something different on step 500:\n",
        "      if i == 500 or i == 1000:\n",
        "        x_curr = x * inv_new_mask\n",
        "      else:\n",
        "        x_curr = x\n",
        "      #x_curr = x\n",
        "\n",
        "      rgb = to_rgb(x_curr)\n",
        "      # Make it blinking with adversaries!\n",
        "      \"\"\"\n",
        "      adv_rgb_color = tf.constant([[[[1.,0., 0.]]]])\n",
        "      t = 0.5 + tf.sin(i/20 / 10) / 2.\n",
        "      adv_mask_period = new_mask * t\n",
        "      img_behind_mask = rgb * new_mask\n",
        "      adv_rgb = adv_rgb_color * adv_mask_period + img_behind_mask * (1 - t)\n",
        "\n",
        "      rgb = rgb * inv_new_mask + adv_rgb\n",
        "      \"\"\"\n",
        "\n",
        "      if num_runs == 1:\n",
        "        rgb = rgb[0]\n",
        "      else:\n",
        "        rgb = tile2d(rgb)#, w=wtile)\n",
        "\n",
        "\n",
        "\n",
        "      vis = zoom(rgb, 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perc virus: {}%\".format(int(exp_percentage_virus * 100))\n",
        "      perturbcolor = (255,0,0)# if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "      vid.add(np.uint8(im))\n",
        "      \n",
        "      if i == 500 or i == 1000:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(50):\n",
        "          vid.add(np.uint8(im))\n",
        "      \n",
        "    x_orig = orig_ca(x)\n",
        "    x_adv = adv_ca(x)\n",
        "    x = x_orig * inv_new_mask + x_adv * new_mask\n",
        "\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bOeOk40qQN7G"
      },
      "outputs": [],
      "source": [
        "#@title viz red 60% saved run\n",
        "import PIL.ImageFont\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "orig_ca = get_model(target_emoji)\n",
        "red_ca = adversarial_ca\n",
        "\n",
        "red_target_img = load_image_from_file(lizard_red_fp)\n",
        "p = 24\n",
        "exp_x0_seed, red_pad_target = generate_padded_xy(p, red_target_img, orig_ca)\n",
        "\n",
        "red_pad_target_mask = tf.cast(red_pad_target[...,3:4] \u003e= 0.1, tf.float32)\n",
        "\n",
        "exp_percentage_virus = 0.60\n",
        "num_runs = 1\n",
        "x =  np.repeat(exp_x0_seed, num_runs, 0)\n",
        "wtile = 3\n",
        "new_mask = random_mask_not_out_target(num_runs, exp_percentage_virus, \n",
        "                                      red_pad_target_mask).numpy()\n",
        "inv_new_mask = 1. - new_mask\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(4001):\n",
        "    if (i\u003c300 and i%5 == 0) or i%20==0:\n",
        "      \n",
        "      # We want to show something different on step 500:\n",
        "      if i == 500:\n",
        "        x_curr = x * inv_new_mask\n",
        "      else:\n",
        "        x_curr = x\n",
        "      #x_curr = x\n",
        "\n",
        "      rgb = to_rgb(x_curr)\n",
        "      # Make it blinking with adversaries!\n",
        "      \"\"\"\n",
        "      adv_rgb_color = tf.constant([[[[1.,0., 0.]]]])\n",
        "      t = 0.5 + tf.sin(i/20 / 10) / 2.\n",
        "      adv_mask_period = new_mask * t\n",
        "      img_behind_mask = rgb * new_mask\n",
        "      adv_rgb = adv_rgb_color * adv_mask_period + img_behind_mask * (1 - t)\n",
        "\n",
        "      rgb = rgb * inv_new_mask + adv_rgb\n",
        "      \"\"\"\n",
        "\n",
        "      if num_runs == 1:\n",
        "        rgb = rgb[0]\n",
        "      else:\n",
        "        rgb = tile2d(rgb, w=wtile)\n",
        "\n",
        "\n",
        "\n",
        "      vis = zoom(rgb, 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perc virus: {}%\".format(int(exp_percentage_virus * 100))\n",
        "      perturbcolor = (255,0,0)# if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "      vid.add(np.uint8(im))\n",
        "       \n",
        "      if i == 500:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(80):\n",
        "          vid.add(np.uint8(im))\n",
        "      \n",
        "    x_orig = orig_ca(x)\n",
        "    x_adv = red_ca(x)\n",
        "    x = x_orig * inv_new_mask + x_adv * new_mask\n",
        "\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1SQtOR3bYhDo"
      },
      "outputs": [],
      "source": [
        "#@title Vanilla run\n",
        "\n",
        "orig_ca = get_model(target_emoji)\n",
        "\n",
        "seed = np.zeros([h, w, CHANNEL_N], np.float32)\n",
        "seed[h//2, w//2, 3:] = 1.0\n",
        "seed = seed[None, ...]\n",
        "x = seed\n",
        "wtile = 3\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(201):\n",
        "    rgb = to_rgb(x)\n",
        "\n",
        "    rgb = rgb[0]\n",
        "\n",
        "    vis = zoom(rgb, 4).clip(0, 1)\n",
        "    #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "    im = np.uint8(vis*255)\n",
        "\n",
        "    im = PIL.Image.fromarray(im)\n",
        "    vid.add(np.uint8(im))\n",
        "    \n",
        "    x = orig_ca(x)\n",
        "\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JaITnQv0k1iY"
      },
      "source": [
        "# Experiment: State perturbations\n",
        "\n",
        "We get a pretrained model (lizard) and train a state mutation to make it grow with missing limbs or different colors.\n",
        "\n",
        "The mutation is applied at all states."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VBw2rdr1k1iZ"
      },
      "outputs": [],
      "source": [
        "target_label = 'lizard_blue' #@param [\"liz_no_tail\", \"liz_no_leg\", \"liz_no_head\", \"lizard_no_arm\", \"lizard_red\", \"lizard_blue\"]\n",
        "\n",
        "target_img = load_image_from_file(target_map[target_label])\n",
        "imshow(zoom(to_rgb(target_img), 2), fmt='png')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l8S9Y-HcVmQh"
      },
      "outputs": [],
      "source": [
        "# compose picture of all perturbations\n",
        "perts_fp_list = [lizard_no_tail_fp, lizard_no_leg_fp, lizard_no_head_fp, lizard_no_arm_fp, lizard_red_fp, lizard_blue_fp]\n",
        "perts_img_list = [to_rgb(load_image_from_file(pfp)) for pfp in perts_fp_list]\n",
        "\n",
        "imshow(zoom(tile2d(perts_img_list), 4))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QAmhVqgik1ia"
      },
      "outputs": [],
      "source": [
        "#@title Initialize Training { vertical-output: true}\n",
        "\n",
        "p = TARGET_PADDING\n",
        "pad_target = tf.pad(target_img, [(p, p), (p, p), (0, 0)])\n",
        "h, w = pad_target.shape[:2]\n",
        "seed = np.zeros([h, w, CHANNEL_N], np.float32)\n",
        "seed[h//2, w//2, 3:] = 1.0\n",
        "\n",
        "def loss_f(x):\n",
        "  return tf.reduce_mean(tf.square(to_rgba(x)-pad_target), [-2, -3, -1])\n",
        "\n",
        "target_emoji = '🦎' # '🦋'\n",
        "ca = get_model(emoji=target_emoji)\n",
        "\n",
        "# Generate the final state to modify.\n",
        "x0_seed = seed[None, ...]\n",
        "for i in tf.range(200):\n",
        "  x0_seed = ca(x0_seed)\n",
        "imshow(zoom(to_rgb(x0_seed[0]), 2), fmt='png')\n",
        "\n",
        "loss_log = []\n",
        "\n",
        "lr = 2e-3\n",
        "lr_sched = tf.keras.optimizers.schedules.PiecewiseConstantDecay(\n",
        "    [2000], [lr, lr*0.1])\n",
        "trainer = tf.keras.optimizers.Adam(lr_sched)\n",
        "\n",
        "loss0 = loss_f(seed).numpy()\n",
        "\n",
        "initial_pool_elems = np.concatenate([\n",
        "    np.repeat(seed[None, ...], POOL_SIZE // 2, 0),\n",
        "    np.repeat(x0_seed, POOL_SIZE // 2, 0)], 0)\n",
        "pool = SamplePool(x=initial_pool_elems)\n",
        "\n",
        "# New mutation tensor. it's a 2d matrix we'll treat as symmetric: [numch, numch]\n",
        "mutation_tensor = tf.Variable(tf.eye(CHANNEL_N))\n",
        "\n",
        "!mkdir -p train_log \u0026\u0026 rm -f train_log/*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m6_asEHrk1ia"
      },
      "outputs": [],
      "source": [
        "#@title Training Loop {vertical-output: true}\n",
        "\n",
        "@tf.function\n",
        "def train_step(x):\n",
        "  iter_n = tf.random.uniform([], 64, 96, tf.int32)\n",
        "  with tf.GradientTape() as g:\n",
        "    m_upper = tf.linalg.band_part(mutation_tensor, 0, -1)\n",
        "    mutation_symm_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(\n",
        "        tf.linalg.diag_part(mutation_tensor))\n",
        "    for i in tf.range(iter_n):\n",
        "      # Apply mutation.\n",
        "      # make sure you don't touch stuff outside.\n",
        "      lm = tf.cast(get_living_mask(x), tf.float32)\n",
        "      x = x @ mutation_symm_t\n",
        "      # saturate mutation past abs(3)\n",
        "      x = tf.clip_by_value(x, -3., 3.)\n",
        "      x *= lm\n",
        "      x = ca(x)\n",
        "    loss = tf.reduce_mean(loss_f(x))\n",
        "  grads = g.gradient(loss, [mutation_tensor])\n",
        "  grads = [g/(tf.norm(g)+1e-8) for g in grads]\n",
        "  trainer.apply_gradients(zip(grads, [mutation_tensor]))\n",
        "  return x, loss\n",
        "\n",
        "for i in range(8000+1):\n",
        "  batch = pool.sample(BATCH_SIZE)\n",
        "  x0 = batch.x\n",
        "  loss_rank = loss_f(x0).numpy().argsort()[::-1]\n",
        "  x0 = x0[loss_rank]\n",
        "  x0[:1] = seed\n",
        "  x0[1:2] = x0_seed[0]\n",
        "  if DAMAGE_N:\n",
        "    damage = 1.0-make_circle_masks(DAMAGE_N, h, w).numpy()[..., None]\n",
        "    x0[-DAMAGE_N:] *= damage\n",
        "\n",
        "  x, loss = train_step(x0)\n",
        "\n",
        "  batch.x[:] = x\n",
        "  batch.commit()\n",
        "\n",
        "  step_i = len(loss_log)\n",
        "  loss_log.append(loss.numpy())\n",
        "  \n",
        "  if step_i%10 == 0:\n",
        "    generate_pool_figures(pool, step_i)\n",
        "  if step_i%100 == 0:\n",
        "    clear_output()\n",
        "    visualize_batch(x0, x, step_i)\n",
        "    plot_loss(loss_log)\n",
        "    export_model(ca, 'train_log/%04d'%step_i)\n",
        "\n",
        "  print('\\r step: %d, log10(loss): %.3f'%(len(loss_log), np.log10(loss)), end='')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "99A96z6lEk6t"
      },
      "outputs": [],
      "source": [
        "# Save the symmetric matrix, in case you want to use combinations of them.\n",
        "# showing an example.\n",
        "m_upper = tf.linalg.band_part(mutation_tensor, 0, -1)\n",
        "#tail_sym_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(tf.linalg.diag_part(mutation_tensor))\n",
        "#leg_sym_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(tf.linalg.diag_part(mutation_tensor))\n",
        "#head_sym_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(tf.linalg.diag_part(mutation_tensor))\n",
        "#arm_sym_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(tf.linalg.diag_part(mutation_tensor))\n",
        "red_sym_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(tf.linalg.diag_part(mutation_tensor))\n",
        "#blue_sym_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(tf.linalg.diag_part(mutation_tensor))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NExpbvw5k1ib"
      },
      "outputs": [],
      "source": [
        "#@title visualize training run\n",
        "EXP = \"default\"\n",
        "import PIL.ImageFont\n",
        "\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "x = seed[None, ...]\n",
        "if EXP == \"default\":\n",
        "  m_upper = tf.linalg.band_part(mutation_tensor, 0, -1)\n",
        "  mutation_symm_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(\n",
        "      tf.linalg.diag_part(mutation_tensor))\n",
        "elif EXP == \"diag\":\n",
        "  # Experiment: only use the diagonal!\n",
        "  mutation_symm_t = tf.linalg.tensor_diag(\n",
        "      tf.linalg.diag_part(mutation_tensor))\n",
        "elif EXP == \"eigvec\":\n",
        "  m_upper = tf.linalg.band_part(mutation_tensor, 0, -1)\n",
        "  mutation_symm_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(\n",
        "      tf.linalg.diag_part(mutation_tensor))\n",
        "\n",
        "  eigval, eigvec = np.linalg.eigh(mutation_symm_t)\n",
        "  eigdelta = eigval - tf.ones([CHANNEL_N])\n",
        "  eigval = tf.ones([CHANNEL_N]) + eigdelta * 0.5\n",
        "  mutation_symm_t = eigvec @ tf.linalg.tensor_diag(eigval) @ np.transpose(eigvec) \n",
        "\n",
        "perturb_step = lambda i: i \u003c 500 or i \u003e= 1000\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(2500):\n",
        "    if i\u003c200 or i%5 == 0:\n",
        "      vis = zoom(to_rgb(x[0]), 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perturbation: {}\".format(\"ON\" if perturb_step(i) else \"OFF\")\n",
        "      perturbcolor = (255,0,0) if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "      vid.add(np.uint8(im))\n",
        "      if i == 500 or i == 1000:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(50):\n",
        "          vid.add(np.uint8(im))\n",
        "    if perturb_step(i):\n",
        "      lm = tf.cast(get_living_mask(x), tf.float32)\n",
        "      x = x @ mutation_symm_t\n",
        "      x = tf.clip_by_value(x, -3., 3.)\n",
        "      x *= lm\n",
        "    x = ca(x)\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zVMjsehQ7jTM"
      },
      "outputs": [],
      "source": [
        "# Create a mosaic of direction mutations!\n",
        "import PIL.ImageFont\n",
        "\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "# Generate 9 mutation matrices:\n",
        "all_mutation_symm_t = []\n",
        "all_distortions = [1.0, 0.5, 0.3, 0.1, 0.0, -0.1, -0.3, -0.5, -1.0]\n",
        "for d in all_distortions:\n",
        "  m_upper = tf.linalg.band_part(mutation_tensor, 0, -1)\n",
        "  mutation_symm_t = m_upper + tf.transpose(m_upper) - tf.linalg.tensor_diag(\n",
        "      tf.linalg.diag_part(mutation_tensor))\n",
        "  \n",
        "  # d*M + (1-d)*I\n",
        "  mutation_symm_t = d * mutation_symm_t + (1. - d) * tf.eye(CHANNEL_N)\n",
        "\n",
        "  all_mutation_symm_t.append(mutation_symm_t)\n",
        "mutation_symm_t = np.stack(all_mutation_symm_t)\n",
        "print(mutation_symm_t.shape)\n",
        "\n",
        "x = np.repeat(seed[None, ...], 9, 0)\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(500):\n",
        "    if i\u003c200 or i%5 == 0:\n",
        "      vis = zoom(tile2d(to_rgb(x), 3), 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perturbation: {}\".format(\"ON\" if perturb_step(i) else \"OFF\")\n",
        "      perturbcolor = (255,0,0) if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "\n",
        "      # Add mutation direction text.\n",
        "      for idx, d in enumerate(all_distortions):\n",
        "        disttext = \"direction: {}\".format(d)\n",
        "        if d == 1.0:\n",
        "          disttext += \" (train config)\"\n",
        "        if d == 0.0:\n",
        "          disttext += \" (NOOP)\"\n",
        "        x_unit, y_unit = im.width // 3, im.height // 3\n",
        "        x_displacement = x_unit // 5\n",
        "        y_displacement = int(y_unit * 0.9)\n",
        "        placement = (x_unit * (idx % 3) + x_displacement,\n",
        "                     y_unit * (idx // 3) + y_displacement)\n",
        "        draw.text(placement,disttext,(0,0,0),font=font)\n",
        "\n",
        "\n",
        "      vid.add(np.uint8(im))\n",
        "      if i == 500 or i == 1000:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(50):\n",
        "          vid.add(np.uint8(im))\n",
        "    if perturb_step(i):\n",
        "      lm = tf.cast(get_living_mask(x), tf.float32)\n",
        "      x_t = []\n",
        "      for d_idx in range(len(all_distortions)):\n",
        "        x_ti = x[d_idx] @ mutation_symm_t[d_idx]\n",
        "        x_ti = tf.clip_by_value(x_ti, -3., 3.)\n",
        "        x_t.append(x_ti)\n",
        "      x = tf.stack(x_t)\n",
        "      x *= lm\n",
        "    x = ca(x)\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "en1-uUZciiM5"
      },
      "outputs": [],
      "source": [
        "# @title Save data for external demo usage\n",
        "\n",
        "perturbations_str = json.dumps(np.stack(\n",
        "    [tail_sym_t, leg_sym_t, head_sym_t, \n",
        "     arm_sym_t, red_sym_t, blue_sym_t]).tolist())\n",
        "\n",
        "emoji = '🦎'\n",
        "model_str = get_model(emoji, output='json')\n",
        "\n",
        "data_js = '''\n",
        "  window.GRAPH_URL = URL.createObjectURL(new Blob([`%s`], {type: 'application/json'}));\n",
        "  window.PERTURBATIONS = %s\n",
        "'''%(model_str, perturbations_str)\n",
        "\n",
        "with open(\"gca_data.js\", \"w\") as f:\n",
        "  f.write(data_js)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2TNYlCDlKC_z"
      },
      "outputs": [],
      "source": [
        "# @title Download trained perturbations from Github\n",
        "!wget -O growing_ca_saved_perturbations.zip 'https://github.com/google-research/self-organising-systems/blob/master/adversarial_reprogramming_ca/assets/growing_ca_saved_perturbations.zip?raw=true'\n",
        "!unzip -oq \"growing_ca_saved_perturbations.zip\" -d \"saved_perturbations\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bLYBtOF2K-jE"
      },
      "outputs": [],
      "source": [
        "# note you can also alternatively just upload the zip file to colab and unzip.\n",
        "#!unzip -oq \"/content/growing_ca_saved_perturbations.zip\" -d \"saved_perturbations\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lecToiJxZjbU"
      },
      "outputs": [],
      "source": [
        "# @title Load saved matrices (for demos and visualizations)\n",
        "\n",
        "load_dir = \"saved_perturbations/\"\n",
        "\n",
        "leg_sym_t = np.load(load_dir + \"leg_sym_t.npy\")\n",
        "tail_sym_t = np.load(load_dir + \"tail_sym_t.npy\")\n",
        "head_sym_t = np.load(load_dir + \"head_sym_t.npy\")\n",
        "arm_sym_t = np.load(load_dir + \"arm_sym_t.npy\")\n",
        "red_sym_t = np.load(load_dir + \"red_sym_t.npy\")\n",
        "blue_sym_t = np.load(load_dir + \"blue_sym_t.npy\")\n",
        "\n",
        "perts_mat_list = [tail_sym_t, leg_sym_t, head_sym_t, arm_sym_t, red_sym_t, blue_sym_t]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f1fctO7h3JNd"
      },
      "outputs": [],
      "source": [
        "#@title TensorFlow.js Demo {run:\"auto\", vertical-output: true}\n",
        "#@markdown Select \"CHECKPOINT\" model to load the checkpoint created by running cells from the \"Training\" section of this notebook\n",
        "import IPython.display\n",
        "\n",
        "emoji = '🦎'\n",
        "\n",
        "perturbations_str = json.dumps(np.stack(\n",
        "    [tail_sym_t, leg_sym_t, head_sym_t, \n",
        "     arm_sym_t, red_sym_t, blue_sym_t]).tolist())\n",
        "\n",
        "#@markdown Shift-click to seed the pattern\n",
        "\n",
        "model_str = get_model(emoji, output='json')\n",
        "\n",
        "data_js = '''\n",
        "  window.GRAPH_URL = URL.createObjectURL(new Blob([`%s`], {type: 'application/json'}));\n",
        "  window.PERTURBATIONS = %s\n",
        "'''%(model_str, perturbations_str)\n",
        "\n",
        "display(IPython.display.Javascript(data_js))\n",
        "\n",
        "\n",
        "IPython.display.HTML('''\n",
        "\u003cscript src=\"https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@2.3.0/dist/tf.min.js\"\u003e\u003c/script\u003e\n",
        "\u003cscript src=\"https://cdnjs.cloudflare.com/ajax/libs/cash/4.1.2/cash.min.js\"\u003e\u003c/script\u003e\n",
        "\n",
        "\u003ccanvas id='canvas' style=\"border: 1px solid black; image-rendering: pixelated;\"\u003e\u003c/canvas\u003e\n",
        "\u003cdiv\u003e\u003cbutton type=\"button\" id=\"reset\"\u003eReset\u003c/button\u003e\u003c/div\u003e\n",
        "\u003cdiv class=\"slidecontainer\"\u003e\n",
        "    tailPerturbation:\n",
        "    \u003cinput type=\"range\" min=\"-1.0\" max=\"1.0\" value=\"0.0\" class=\"slider\" \n",
        "      step=\"0.05\" id=\"tailSlider\"\u003e\n",
        "    \u003cspan id='tailPerturbation'\u003e0.0\u003c/span\u003e\n",
        "\u003c/div\u003e\n",
        "\u003cdiv class=\"slidecontainer\"\u003e\n",
        "    legPerturbation:\n",
        "    \u003cinput type=\"range\" min=\"-1.0\" max=\"1.0\" value=\"0.0\" class=\"slider\" \n",
        "      step=\"0.05\" id=\"legSlider\"\u003e\n",
        "    \u003cspan id='legPerturbation'\u003e0.0\u003c/span\u003e\n",
        "\u003c/div\u003e\n",
        "\u003cdiv class=\"slidecontainer\"\u003e\n",
        "    headPerturbation:\n",
        "    \u003cinput type=\"range\" min=\"-1.0\" max=\"1.0\" value=\"0.0\" class=\"slider\" \n",
        "      step=\"0.05\" id=\"headSlider\"\u003e\n",
        "    \u003cspan id='headPerturbation'\u003e0.0\u003c/span\u003e\n",
        "\u003c/div\u003e\n",
        "\u003cdiv class=\"slidecontainer\"\u003e\n",
        "    armPerturbation:\n",
        "    \u003cinput type=\"range\" min=\"-1.0\" max=\"1.0\" value=\"0.0\" class=\"slider\" \n",
        "      step=\"0.05\" id=\"armSlider\"\u003e\n",
        "    \u003cspan id='armPerturbation'\u003e0.0\u003c/span\u003e\n",
        "\u003c/div\u003e\n",
        "\u003cdiv class=\"slidecontainer\"\u003e\n",
        "    redPerturbation:\n",
        "    \u003cinput type=\"range\" min=\"-1.0\" max=\"1.0\" value=\"0.0\" class=\"slider\" \n",
        "      step=\"0.05\" id=\"redSlider\"\u003e\n",
        "    \u003cspan id='redPerturbation'\u003e0.0\u003c/span\u003e\n",
        "\u003c/div\u003e\n",
        "\u003cdiv class=\"slidecontainer\"\u003e\n",
        "    bluePerturbation:\n",
        "    \u003cinput type=\"range\" min=\"-1.0\" max=\"1.0\" value=\"0.0\" class=\"slider\" \n",
        "      step=\"0.05\" id=\"blueSlider\"\u003e\n",
        "    \u003cspan id='bluePerturbation'\u003e0.0\u003c/span\u003e\n",
        "\u003c/div\u003e\n",
        "\u003cinput type=\"checkbox\" id=\"forcesum1\" name=\"forcesum1\"\u003e\n",
        "\u003clabel for=\"forcesum1\"\u003e Force sum leq 1\u003c/label\u003e\u003cbr\u003e\n",
        "\u003cscript\u003e\n",
        "  \"use strict\";\n",
        "  \n",
        "  const sleep = (ms)=\u003enew Promise(resolve =\u003e setTimeout(resolve, ms));\n",
        "  \n",
        "  const parseConsts = model_graph=\u003e{\n",
        "    const dtypes = {'DT_INT32':['int32', 'intVal', Int32Array],\n",
        "                    'DT_FLOAT':['float32', 'floatVal', Float32Array]};\n",
        "    \n",
        "    const consts = {};\n",
        "    model_graph.modelTopology.node.filter(n=\u003en.op=='Const').forEach((node=\u003e{\n",
        "      const v = node.attr.value.tensor;\n",
        "      const [dtype, field, arrayType] = dtypes[v.dtype];\n",
        "      if (!v.tensorShape.dim) {\n",
        "        consts[node.name] = [tf.scalar(v[field][0], dtype)];\n",
        "      } else {\n",
        "        const shape = v.tensorShape.dim.map(d=\u003eparseInt(d.size));\n",
        "        let arr;\n",
        "        if (v.tensorContent) {\n",
        "          const data = atob(v.tensorContent);\n",
        "          const buf = new Uint8Array(data.length);\n",
        "          for (var i=0; i\u003cdata.length; ++i) {\n",
        "            buf[i] = data.charCodeAt(i);\n",
        "          }\n",
        "          arr = new arrayType(buf.buffer);\n",
        "        } else {\n",
        "          const size = shape.reduce((a, b)=\u003ea*b);\n",
        "          arr = new arrayType(size);\n",
        "          arr.fill(v[field][0]);\n",
        "        }\n",
        "        consts[node.name] = [tf.tensor(arr, shape, dtype)];\n",
        "      }\n",
        "    }));\n",
        "    return consts;\n",
        "  }\n",
        "\n",
        "  let kTail = 0.0;\n",
        "  let kLeg = 0.0;\n",
        "  let kHead = 0.0;\n",
        "  let kArm = 0.0;\n",
        "  let kRed = 0.0;\n",
        "  let kBlue = 0.0;\n",
        "\n",
        "  let forcesum1Ckbx = document.getElementById(\"forcesum1\");\n",
        "  let tailSlider = document.getElementById(\"tailSlider\");\n",
        "  let legSlider = document.getElementById(\"legSlider\");\n",
        "  let headSlider = document.getElementById(\"headSlider\");\n",
        "  let armSlider = document.getElementById(\"armSlider\");\n",
        "  let redSlider = document.getElementById(\"redSlider\");\n",
        "  let blueSlider = document.getElementById(\"blueSlider\");\n",
        "\n",
        "  $('#tailSlider').on('input', e=\u003e{\n",
        "      updateK(\"tail\", parseFloat(e.target.value));\n",
        "  });\n",
        "  $('#legSlider').on('input', e=\u003e{\n",
        "      updateK(\"leg\", parseFloat(e.target.value));\n",
        "  });\n",
        "  $('#headSlider').on('input', e=\u003e{\n",
        "      updateK(\"head\", parseFloat(e.target.value));\n",
        "  });\n",
        "  $('#armSlider').on('input', e=\u003e{\n",
        "      updateK(\"arm\", parseFloat(e.target.value));\n",
        "  });\n",
        "  $('#redSlider').on('input', e=\u003e{\n",
        "      updateK(\"red\", parseFloat(e.target.value));\n",
        "  });\n",
        "  $('#blueSlider').on('input', e=\u003e{\n",
        "      updateK(\"blue\", parseFloat(e.target.value));\n",
        "  });\n",
        "\n",
        "  const updateKUnchecked = (kid, v) =\u003e {\n",
        "    if (kid == \"tail\"){\n",
        "      kTail = v;\n",
        "      $('#tailPerturbation').text(kTail);\n",
        "    } else if (kid == \"leg\"){\n",
        "      kLeg = v;\n",
        "      $('#legPerturbation').text(kLeg);\n",
        "    } else if (kid == \"head\"){\n",
        "      kHead = v;\n",
        "      $('#headPerturbation').text(kHead);\n",
        "    } else if (kid == \"arm\"){\n",
        "      kArm = v;\n",
        "      $('#armPerturbation').text(kArm);\n",
        "    } else if (kid == \"red\"){\n",
        "      kRed = v;\n",
        "      $('#redPerturbation').text(kRed);\n",
        "    } else if (kid == \"blue\"){\n",
        "      kBlue = v;\n",
        "      $('#bluePerturbation').text(kBlue);\n",
        "    } else {\n",
        "      console.log(\"ERROR!\");\n",
        "    }\n",
        "  }\n",
        "\n",
        "  const updateK = (kid, v) =\u003e {\n",
        "      if (forcesum1Ckbx.checked == false) {\n",
        "        updateKUnchecked(kid, v);\n",
        "      } else {\n",
        "        // You cannot go over 1.\n",
        "        let vAbs = Math.abs(v);\n",
        "        const vSign = Math.sign(v);\n",
        "        const kTailAbs = Math.abs(kTail);\n",
        "        const kTailSign = Math.sign(kTail);\n",
        "        const kLegAbs = Math.abs(kLeg);\n",
        "        const kLegSign = Math.sign(kLeg);\n",
        "        const kHeadAbs = Math.abs(kHead);\n",
        "        const kHeadSign = Math.sign(kHead);\n",
        "        const kArmAbs = Math.abs(kArm);\n",
        "        const kArmSign = Math.sign(kArm);\n",
        "        const kRedAbs = Math.abs(kRed);\n",
        "        const kRedSign = Math.sign(kRed);\n",
        "        const kBlueAbs = Math.abs(kBlue);\n",
        "        const kBlueSign = Math.sign(kBlue);\n",
        "\n",
        "        let kCurrAbs;\n",
        "        if (kid == \"tail\"){\n",
        "          kCurrAbs = kTailAbs;\n",
        "        } else if (kid == \"leg\"){\n",
        "          kCurrAbs = kLegAbs;\n",
        "        } else if (kid == \"head\"){\n",
        "          kCurrAbs = kHeadAbs;\n",
        "        } else if (kid == \"arm\"){\n",
        "          kCurrAbs = kArmAbs;\n",
        "        } else if (kid == \"red\"){\n",
        "          kCurrAbs = kRedAbs;\n",
        "        } else if (kid == \"blue\"){\n",
        "          kCurrAbs = kBlueAbs;\n",
        "        }\n",
        "        let totK = vAbs + kTailAbs + kLegAbs + kHeadAbs + kArmAbs\n",
        "              + kRedAbs + kBlueAbs - kCurrAbs;\n",
        "        if (totK \u003c= 1.0) {\n",
        "          // No problem here, do just like you did before.\n",
        "          updateKUnchecked(kid, v);\n",
        "        } else {\n",
        "          // Prevent v from going over 1.\n",
        "          if (vAbs \u003e 1.0) {\n",
        "            vAbs = 1.0;\n",
        "            totK = vAbs + kTailAbs + kLegAbs + kHeadAbs + kArmAbs\n",
        "              + kRedAbs + kBlueAbs - kCurrAbs;\n",
        "          }\n",
        "          // Subtract the excess from the rest.\n",
        "          const excess = totK - 1.0;\n",
        "\n",
        "          const tailContrib = kid == \"tail\" ? 0.0 : kTailAbs;\n",
        "          const legContrib = kid == \"leg\" ? 0.0 : kLegAbs;\n",
        "          const headContrib = kid == \"head\" ? 0.0 : kHeadAbs;\n",
        "          const armContrib = kid == \"arm\" ? 0.0 : kArmAbs;\n",
        "          const redContrib = kid == \"red\" ? 0.0 : kRedAbs;\n",
        "          const blueContrib = kid == \"blue\" ? 0.0 : kBlueAbs;\n",
        "          const totContrib = tailContrib + legContrib + headContrib +\n",
        "              armContrib + redContrib + blueContrib;\n",
        "\n",
        "          let tailDecr = 0.0;\n",
        "          let legDecr = 0.0;\n",
        "          let headDecr = 0.0;\n",
        "          let armDecr = 0.0;\n",
        "          let redDecr = 0.0;\n",
        "          let blueDecr = 0.0;\n",
        "          if (totContrib \u003e 1e-6) {\n",
        "            tailDecr = tailContrib / totContrib * excess;\n",
        "            legDecr = legContrib / totContrib * excess;\n",
        "            headDecr = headContrib / totContrib * excess;\n",
        "            armDecr = armContrib / totContrib * excess;\n",
        "            redDecr = redContrib / totContrib * excess;\n",
        "            blueDecr = blueContrib / totContrib * excess;\n",
        "          }\n",
        "\n",
        "          kTail = kid == \"tail\" ? vAbs * vSign : kTailSign * (kTailAbs - tailDecr);\n",
        "          kLeg = kid == \"leg\" ? vAbs * vSign : kLegSign * (kLegAbs - legDecr);\n",
        "          kHead = kid == \"head\" ? vAbs * vSign : kHeadSign * (kHeadAbs - headDecr);\n",
        "          kArm = kid == \"arm\" ? vAbs * vSign : kArmSign * (kArmAbs - armDecr);\n",
        "          kRed = kid == \"red\" ? vAbs * vSign : kRedSign * (kRedAbs - redDecr);\n",
        "          kBlue = kid == \"blue\" ? vAbs * vSign : kBlueSign * (kBlueAbs - blueDecr);\n",
        "          $('#tailPerturbation').text(kTail);\n",
        "          tailSlider.value = kTail;\n",
        "          $('#legPerturbation').text(kLeg);\n",
        "          legSlider.value = kLeg;\n",
        "          $('#headPerturbation').text(kHead);\n",
        "          headSlider.value = kHead;\n",
        "          $('#armPerturbation').text(kArm);\n",
        "          armSlider.value = kArm;\n",
        "          $('#redPerturbation').text(kRed);\n",
        "          redSlider.value = kRed;\n",
        "          $('#bluePerturbation').text(kBlue);\n",
        "          blueSlider.value = kBlue;\n",
        "        }\n",
        "      }\n",
        "      updatePerturbation();\n",
        "  }\n",
        "\n",
        "\n",
        "  let perturbationMatrix = tf.eye(16);\n",
        "\n",
        "  const perturbations = tf.tensor(PERTURBATIONS);\n",
        "  const tailPertM = perturbations.gather([0]).squeeze();\n",
        "  const legPertM = perturbations.gather([1]).squeeze();\n",
        "  const headPertM = perturbations.gather([2]).squeeze();\n",
        "  const armPertM = perturbations.gather([3]).squeeze();\n",
        "  const redPertM = perturbations.gather([4]).squeeze();\n",
        "  const bluePertM = perturbations.gather([5]).squeeze();\n",
        "\n",
        "  const I = tf.eye(16);\n",
        "  const updatePerturbation = () =\u003e {\n",
        "      let kI = 1.0 - kTail - kLeg - kHead - kArm - kRed - kBlue;\n",
        "      perturbationMatrix = I.mul(kI).\n",
        "        add(tailPertM.mul(kTail)).\n",
        "        add(legPertM.mul(kLeg)).\n",
        "        add(headPertM.mul(kHead)).\n",
        "        add(armPertM.mul(kArm)).\n",
        "        add(redPertM.mul(kRed)).\n",
        "        add(bluePertM.mul(kBlue));\n",
        "  }\n",
        "\n",
        "  const run = async ()=\u003e{\n",
        "    const r = await fetch(GRAPH_URL);\n",
        "    const consts = parseConsts(await r.json());\n",
        "    \n",
        "    const model = await tf.loadGraphModel(GRAPH_URL);\n",
        "    Object.assign(model.weights, consts);\n",
        "\n",
        "    \n",
        "    let seed = new Array(16).fill(0).map((x, i)=\u003ei\u003c3?0:1);\n",
        "    seed = tf.tensor(seed, [1, 1, 1, 16]);\n",
        "    \n",
        "    const D = 96;\n",
        "    const initState = tf.tidy(()=\u003e{\n",
        "      const D2 = D/2;\n",
        "      const a = seed.pad([[0, 0], [D2-1, D2], [D2-1, D2], [0,0]]);\n",
        "      return a;\n",
        "    });\n",
        "    \n",
        "    const state = tf.variable(initState);\n",
        "    const [_, h, w, ch] = state.shape;\n",
        "    \n",
        "\n",
        "    $('#reset').on('click', e=\u003e{\n",
        "        tf.tidy(()=\u003e{\n",
        "          state.assign(initState);\n",
        "        });\n",
        "    });\n",
        "\n",
        "\n",
        "    const damage = (x, y, r)=\u003e{\n",
        "      tf.tidy(()=\u003e{\n",
        "        const rx = tf.range(0, w).sub(x).div(r).square().expandDims(0);\n",
        "        const ry = tf.range(0, h).sub(y).div(r).square().expandDims(1);\n",
        "        const mask = rx.add(ry).greater(1.0).expandDims(2);\n",
        "        state.assign(state.mul(mask));\n",
        "      });\n",
        "    }\n",
        "    \n",
        "    const plantSeed = (x, y)=\u003e{\n",
        "      const x2 = w-x-seed.shape[2];\n",
        "      const y2 = h-y-seed.shape[1];\n",
        "      if (x\u003c0 || x2\u003c0 || y2\u003c0 || y2\u003c0)\n",
        "        return;\n",
        "      tf.tidy(()=\u003e{\n",
        "        const a = seed.pad([[0, 0], [y, y2], [x, x2], [0,0]]);\n",
        "        state.assign(state.add(a));\n",
        "      });\n",
        "    }\n",
        "    \n",
        "    const scale = 4;\n",
        "    \n",
        "    const canvas = document.getElementById('canvas');\n",
        "    const ctx = canvas.getContext('2d');\n",
        "    canvas.width = w;\n",
        "    canvas.height = h;\n",
        "    canvas.style.width = `${w*scale}px`;\n",
        "    canvas.style.height = `${h*scale}px`;\n",
        "    \n",
        "    canvas.onmousedown = e=\u003e{\n",
        "      const x = Math.floor(e.clientX/scale);\n",
        "        const y = Math.floor(e.clientY/scale);\n",
        "        if (e.buttons == 1) {\n",
        "          if (e.shiftKey) {\n",
        "            plantSeed(x, y);  \n",
        "          } else {\n",
        "            damage(x, y, 8);\n",
        "          }\n",
        "        }\n",
        "    }\n",
        "    canvas.onmousemove = e=\u003e{\n",
        "      const x = Math.floor(e.clientX/scale);\n",
        "      const y = Math.floor(e.clientY/scale);\n",
        "      if (e.buttons == 1 \u0026\u0026 !e.shiftKey) {\n",
        "        damage(x, y, 8);\n",
        "      }\n",
        "    }\n",
        "\n",
        "    function step() {\n",
        "      tf.tidy(()=\u003e{\n",
        "        let new_state = model.execute(\n",
        "            {x:state, fire_rate:tf.tensor(0.5),\n",
        "            angle:tf.tensor(0.0), step_size:tf.tensor(1.0)}, ['Identity']);\n",
        "        new_state = new_state.reshape([-1, 16]);\n",
        "        new_state = new_state.matMul(perturbationMatrix).reshape([1, D, D, 16]);\n",
        "        new_state = new_state.clipByValue(-3., +3.);\n",
        "        state.assign(new_state);\n",
        "      });\n",
        "    }\n",
        "\n",
        "    function render() {\n",
        "      step();\n",
        "\n",
        "      const imageData = tf.tidy(()=\u003e{\n",
        "        const rgba = state.slice([0, 0, 0, 0], [-1, -1, -1, 4]);\n",
        "        const a = state.slice([0, 0, 0, 3], [-1, -1, -1, 1]);\n",
        "        const img = tf.tensor(1.0).sub(a).add(rgba).mul(255);\n",
        "        const rgbaBytes = new Uint8ClampedArray(img.dataSync());\n",
        "        return new ImageData(rgbaBytes, w, h);\n",
        "      });\n",
        "      ctx.putImageData(imageData, 0, 0);\n",
        "\n",
        "      requestAnimationFrame(render);\n",
        "    }\n",
        "    render();\n",
        "  }\n",
        "  run();\n",
        "  \n",
        "\u003c/script\u003e\n",
        "''')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IVFNqnhKXUaF"
      },
      "source": [
        "## Make videos for article"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XUXrPFdpXb1t"
      },
      "outputs": [],
      "source": [
        "#@title visualize runs\n",
        "import PIL.ImageFont\n",
        "\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "\n",
        "perturb_step = lambda i: i \u003c 500 or i \u003e= 1000\n",
        "\n",
        "runs_per_perturb = []\n",
        "for sym_t in perts_mat_list:\n",
        "  x = seed[None, ...]\n",
        "\n",
        "  run_frames = []\n",
        "  run_texts = []\n",
        "  for i in tqdm.trange(1500):\n",
        "    if i\u003c200 or i%5 == 0:\n",
        "      vis = zoom(to_rgb(x[0]), 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "      run_frames.append(im)\n",
        "\n",
        "      # Save text for rendering.\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perturbation: {}\".format(\"ON\" if perturb_step(i) else \"OFF\")\n",
        "      perturbcolor = (255,0,0) if perturb_step(i) else (0,255,0)\n",
        "      run_texts.append((steptext, perturbtext,perturbcolor))\n",
        "      if i == 500 or i == 1000:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(50):\n",
        "          run_frames.append(im)\n",
        "          run_texts.append((steptext, perturbtext,perturbcolor))\n",
        "    if perturb_step(i):\n",
        "      lm = tf.cast(get_living_mask(x), tf.float32)\n",
        "      x = x @ sym_t\n",
        "      x = tf.clip_by_value(x, -3., 3.)\n",
        "      x *= lm\n",
        "    x = ca(x)\n",
        "\n",
        "  runs_per_perturb.append(run_frames)\n",
        "\n",
        "# generate text for each frame with a skeleton like the one above:\n",
        "\n",
        "\n",
        "num_frames = len(runs_per_perturb[0])\n",
        "\n",
        "fn = 'mosaic_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(num_frames):\n",
        "    frame_images = [run_frames[i] for run_frames in runs_per_perturb]\n",
        "    im = tile2d(frame_images)\n",
        "  \n",
        "    im = PIL.Image.fromarray(im)\n",
        "    draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "    # compute text:\n",
        "    steptext, perturbtext, perturbcolor = run_texts[i]\n",
        "    # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "    draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "    draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "    vid.add(np.uint8(im))\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ipIs1AbvZHq-"
      },
      "outputs": [],
      "source": [
        "#@title interpolating directions\n",
        "import PIL.ImageFont\n",
        "\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "perturb_step = lambda i: True\n",
        "\n",
        "# Generate 5 mutation matrices:\n",
        "pert_1 = tail_sym_t\n",
        "pert_2 = leg_sym_t\n",
        "all_mutation_symm_t = []\n",
        "all_distortions = [1.0, 0.75, 0.5, 0.25, 0.0]\n",
        "for d in all_distortions:\n",
        "  \n",
        "  mutation_symm_t = d * pert_1 + (1. - d) * pert_2\n",
        "\n",
        "\n",
        "  all_mutation_symm_t.append(mutation_symm_t)\n",
        "mutation_symm_t = np.stack(all_mutation_symm_t)\n",
        "print(mutation_symm_t.shape)\n",
        "\n",
        "x = np.repeat(seed[None, ...], 5, 0)\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(1500):\n",
        "    if i\u003c200 or i%5 == 0:\n",
        "      vis = zoom(tile2d(to_rgb(x), 5), 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perturbation: {}\".format(\"ON\" if perturb_step(i) else \"OFF\")\n",
        "      perturbcolor = (255,0,0) if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "\n",
        "      # Add mutation direction text.\n",
        "      for idx, d in enumerate(all_distortions):\n",
        "        disttext = \"tail k: {} | leg k: {}\".format(d, 1.0 - d)\n",
        "        x_unit, y_unit = im.width // 5, im.height // 1\n",
        "        x_displacement = x_unit // 5\n",
        "        y_displacement = int(y_unit * 0.9)\n",
        "        placement = (x_unit * (idx % 5) + x_displacement,\n",
        "                     y_unit * (idx // 5) + y_displacement)\n",
        "        draw.text(placement,disttext,(0,0,0),font=font)\n",
        "\n",
        "\n",
        "      vid.add(np.uint8(im))\n",
        "    if perturb_step(i):\n",
        "      lm = tf.cast(get_living_mask(x), tf.float32)\n",
        "      # Uncertain if there is a faster way:\n",
        "      x_t = []\n",
        "      for d_idx in range(len(all_distortions)):\n",
        "        x_ti = x[d_idx] @ mutation_symm_t[d_idx]\n",
        "        x_ti = tf.clip_by_value(x_ti, -3., 3.)\n",
        "        x_t.append(x_ti)\n",
        "      x = tf.stack(x_t)\n",
        "      x *= lm\n",
        "    x = ca(x)\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bLDquHTaC3_T"
      },
      "outputs": [],
      "source": [
        "# @title combinations of mutations\n",
        "import PIL.ImageFont\n",
        "\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "x = seed[None, ...]\n",
        "\n",
        "mutation_symm_t = tail_sym_t + leg_sym_t - tf.eye(16)\n",
        "\n",
        "perturb_step = lambda i: i \u003c 500 or i \u003e= 1000\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(1500):\n",
        "    if i\u003c200 or i%5 == 0:\n",
        "      vis = zoom(to_rgb(x[0]), 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perturbation: {}\".format(\"ON\" if perturb_step(i) else \"OFF\")\n",
        "      perturbcolor = (255,0,0) if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "      vid.add(np.uint8(im))\n",
        "      if i == 500 or i == 1000:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(50):\n",
        "          vid.add(np.uint8(im))\n",
        "    if perturb_step(i):\n",
        "      lm = tf.cast(get_living_mask(x), tf.float32)\n",
        "      x = x @ mutation_symm_t\n",
        "      x = tf.clip_by_value(x, -3., 3.)\n",
        "      x *= lm\n",
        "    x = ca(x)\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DrWjAyLKIdNQ"
      },
      "outputs": [],
      "source": [
        "# @title Create a mosaic of direction mutations!\n",
        "import PIL.ImageFont\n",
        "\n",
        "from matplotlib import font_manager as fm\n",
        "font_fn = fm.findfont(fm.FontProperties())\n",
        "font = PIL.ImageFont.truetype(font_fn, 16)\n",
        "\n",
        "# Generate 9 mutation matrices:\n",
        "all_mutation_symm_t = []\n",
        "all_distortions = [1.0, 0.5, 0.3, 0.1, 0.0, -0.1, -0.3, -0.5, -1.0]\n",
        "for d in all_distortions:\n",
        "  \n",
        "  # Luca's faster way: d*M + (1-d)*I\n",
        "  mutation_symm_t = d * tail_sym_t + (1. - d) * tf.eye(CHANNEL_N)\n",
        "\n",
        "  \"\"\"\n",
        "  eigval, eigvec = np.linalg.eigh(mutation_symm_t)\n",
        "  eigdelta = eigval - tf.ones([CHANNEL_N])\n",
        "  eigval = tf.ones([CHANNEL_N]) + eigdelta * d\n",
        "  mutation_symm_t = eigvec @ tf.linalg.tensor_diag(eigval) @ np.transpose(eigvec)\n",
        "  \"\"\"\n",
        "\n",
        "  all_mutation_symm_t.append(mutation_symm_t)\n",
        "mutation_symm_t = np.stack(all_mutation_symm_t)\n",
        "print(mutation_symm_t.shape)\n",
        "\n",
        "x = np.repeat(seed[None, ...], 9, 0)\n",
        "\n",
        "fn = 'example_run.mp4'\n",
        "with VideoWriter(fn) as vid:\n",
        "  for i in tqdm.trange(500):\n",
        "    if i\u003c200 or i%5 == 0:\n",
        "      vis = zoom(tile2d(to_rgb(x), 3), 4).clip(0, 1)\n",
        "      #vis = np.concatenate((vis, np.ones((164, vis.shape[1], 3))), axis=0) \n",
        "\n",
        "      im = np.uint8(vis*255)\n",
        "\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      \n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "\n",
        "      # compute text:\n",
        "      steptext = \"Step: {}\".format(i)\n",
        "      perturbtext = \"Perturbation: {}\".format(\"ON\" if perturb_step(i) else \"OFF\")\n",
        "      perturbcolor = (255,0,0) if perturb_step(i) else (0,255,0)\n",
        "      # draw.text((x, y),\"Sample Text\",(r,g,b))\n",
        "      draw.text((0, 0),steptext,(0,0,0),font=font)\n",
        "      draw.text((0, 20),perturbtext,perturbcolor,font=font)\n",
        "\n",
        "      # Add mutation direction text.\n",
        "      for idx, d in enumerate(all_distortions):\n",
        "        disttext = \"direction: {}\".format(d)\n",
        "        if d == 1.0:\n",
        "          disttext += \" (train config)\"\n",
        "        if d == 0.0:\n",
        "          disttext += \" (NOOP)\"\n",
        "        x_unit, y_unit = im.width // 3, im.height // 3\n",
        "        x_displacement = x_unit // 5\n",
        "        y_displacement = int(y_unit * 0.9)\n",
        "        placement = (x_unit * (idx % 3) + x_displacement,\n",
        "                     y_unit * (idx // 3) + y_displacement)\n",
        "        draw.text(placement,disttext,(0,0,0),font=font)\n",
        "\n",
        "\n",
        "      vid.add(np.uint8(im))\n",
        "      if i == 500 or i == 1000:\n",
        "        # add many frames to effectively pause the video.\n",
        "        for _ in range(50):\n",
        "          vid.add(np.uint8(im))\n",
        "    if perturb_step(i):\n",
        "      lm = tf.cast(get_living_mask(x), tf.float32)\n",
        "      # Uncertain if there is a faster way:\n",
        "      x_t = []\n",
        "      for d_idx in range(len(all_distortions)):\n",
        "        x_ti = x[d_idx] @ mutation_symm_t[d_idx]\n",
        "        x_ti = tf.clip_by_value(x_ti, -3., 3.)\n",
        "        x_t.append(x_ti)\n",
        "      x = tf.stack(x_t)\n",
        "      x *= lm\n",
        "    x = ca(x)\n",
        "\n",
        "mvp.ipython_display(fn, loop=True)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "Adversarial Reprogramming of Growing Neural Cellular Automata",
      "provenance": [
        {
          "file_id": "1N2V8auh_Hsl6kL4FOK18_OS4BakU1n4t",
          "timestamp": 1616068450257
        },
        {
          "file_id": "1W9C0wUAajw3n2WdMeeSiCxS2MjDfS_UX",
          "timestamp": 1606746533474
        },
        {
          "file_id": "1u3Tk9L0yqaBPg-kvYkncHtatXakZ-8Gu",
          "timestamp": 1600847640721
        },
        {
          "file_id": "https://github.com/google-research/self-organising-systems/blob/master/notebooks/growing_ca.ipynb",
          "timestamp": 1598861702217
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
